from .myseq import MySeq
from typing import Tuple

class SubstMatrix(object):
    def __init__(self) -> None:
        self.alphabet = list()
        self.sm = dict()

    def build_from_dict(self,dic:dict):
        self.sm = dic
        self.alphabet = list({key[0] for key in dic.keys()})
        return self
    
    def create_submat(self,match:int,mismatch:int,alphabet:str):
        self.alphabet = list(alphabet)
        for c1 in alphabet:
            for c2 in alphabet:
                self.sm[c1+c2] = match if c1 == c2 else mismatch
        return self

    def score_pair(self,c1,c2) -> int:
        if c1 in self.alphabet and c2 in self.alphabet:
            return self.sm[c1+c2]
        else:
            ValueError('char not in alphabet')

    def __getitem__(self,ft) -> int:
        from_word,to_word = ft
        return self.score_pair(from_word,to_word)

    def read_submat_file(self):
        pass
    

class Cell(object):
    def __init__(self,*location) -> None:
        self.ij = location
        self.value = 0
        self.trace = [0]
        self.precursor = []
        self.successor = []
    
    def reset(self) -> None:
        self.value = 0
        self.trace = [0]
        self.precursor = []
        self.successor = []


class Table(object):
    def __init__(self,row:str,colum:str) -> None:
        self.row = '-'+row
        self.colum = '-'+colum
        self.shape = (len(self.row),len(self.colum))
        self.table = [[Cell(i,j) for j in range(self.shape[1])
                                ]for i in range(self.shape[0])]

    def __getitem__(self,ij) -> "Cell":
        i,j = ij
        return self.table[i][j]
    
    def fill_gap(self,gap) -> None:
        for i in range(1,self.shape[0]):
            self[i,0].value = self[i-1,0].value + gap
            self.update_pointer((i-1,0),(i,0),trace_type=3)
        for j in range(1,self.shape[1]):
            self[0,j].value = self[0,j-1].value + gap
            self.update_pointer((0,j-1),(0,j),trace_type=2)
    
    def __clean_gap(self) -> None:
        for i in range(1,self.shape[0]):
            self[i,0].reset()
        for j in range(1,self.shape[1]):
            self[0,j].reset()
    
    def update_pointer(self,fij,tij,trace_type=0) -> None:
        self[tij[0],tij[1]].precursor.append(self[fij[0],fij[1]])
        self[fij[0],fij[1]].successor.append(self[tij[0],tij[1]])
        if trace_type:
            self[tij[0],tij[1]].trace.append(trace_type)
    
    def get_score(self) -> list[list[int]]:
        return [[self[i,j].value for j in range(self.shape[1])]for i in range(self.shape[0])]
    
    def get_trace_back(self) -> list[list[int]]:
        return [[self[i,j].trace[-1] for j in range(self.shape[1])]for i in range(self.shape[0])]

    def needleman_Wunsch(self,sm:"SubstMatrix",gap:int) -> None:
        row_size,colum_size = self.shape
        self.fill_gap(gap)
        for suc_word,i in zip(self.row[1:],range(1,row_size)):
            for tag_word,j in zip(self.colum[1:],range(1,colum_size)):
                index,score = max(enumerate((
                    self[i-1,j].value + gap,
                    self[i,j-1].value + gap,
                    self[i-1,j-1].value + sm[suc_word,tag_word],
                )),key=lambda v: v[1])
                self[i,j].value = score
                # print(f'\t{score}',end='')
                if index == 0:
                    self.update_pointer((i-1,j),(i,j),trace_type=3)
                elif index == 1:
                    self.update_pointer((i,j-1),(i,j),trace_type=2)
                elif index == 2:
                    self.update_pointer((i-1,j-1),(i,j),trace_type=1)
                else:
                    raise ValueError('check max')
            # print()
    
    def smith_Waterman(self,sm:"SubstMatrix",gap:int) -> Tuple["Cell"]:
        self.__clean_gap()
        row_size,colum_size = self.shape
        max_cells = [self[0,0]]
        for suc_word,i in zip(self.row[1:],range(1,row_size)):
            for tag_word,j in zip(self.colum[1:],range(1,colum_size)):
                index,score = max(enumerate((
                    self[i-1,j].value + gap,
                    self[i,j-1].value + gap,
                    self[i-1,j-1].value + sm[suc_word,tag_word],
                    0
                )),key=lambda v: v[1])
                self[i,j].value = score
                # print(f'\t{score}',end='')
                if index == 0:
                    self.update_pointer((i-1,j),(i,j),trace_type=3)
                elif index == 1:
                    self.update_pointer((i,j-1),(i,j),trace_type=2)
                elif index == 2:
                    self.update_pointer((i-1,j-1),(i,j),trace_type=1)
                elif index == 3:
                    pass
                else:
                    raise ValueError('check max')
                
                if self[i,j].value > max_cells[0].value:
                    max_cells = [self[i,j]]
                elif self[i,j].value == max_cells[0].value:
                    max_cells.append(self[i,j])
                else:
                    pass
            # print()
        self.max_cells = max_cells
        return max_cells
    
    def trace_back(self,start_cell:"Cell") -> list[str]:
        current_cell = start_cell
        seq1,seq2 = '',''
        while True:
            trace = current_cell.trace[-1]
            i,j = current_cell.ij
            if trace == 1:
                seq1 = self.row[i] + seq1
                seq2 = self.colum[j] + seq2
            elif trace == 2:
                seq1 = '-' + seq1
                seq2 = self.colum[j] + seq2
            elif trace == 3:
                seq1 = self.row[i] + seq1
                seq2 = '-' + seq2
            
            if not trace:
                break
            else:
                current_cell = current_cell.precursor[0]
            
        return [seq1,seq2]

    def recover_align(self):
        return self.trace_back(self[-1,-1])

    def recover_align_local(self) -> list[list[str]]:
        result = list()
        for start_cell in self.max_cells:
            result.append(self.trace_back(start_cell))
        return result

        

def needleman_Wunsch(seq1:str, seq2:str, sm:dict, gap:int):
    table = Table(seq1,seq2)
    sm = SubstMatrix().build_from_dict(sm)
    table.needleman_Wunsch(sm,gap)
    return table.get_score(),table.get_trace_back()

def smith_Waterman(seq1:str, seq2:str, sm:dict, gap:int):
    table = Table(seq1,seq2)
    sm = SubstMatrix().build_from_dict(sm)
    max_cell = table.smith_Waterman(sm,gap)
    return table.get_score(),table.get_trace_back(),max_cell[-1].value


if __name__ == "__main__":

    seq1 = "ACTAGC"
    seq2 = "ATAGAC"
    gap = -1

    sm = SubstMatrix()
    sm.create_submat(1,-1,MySeq.ALPHABET["DNA"])
    
    aligner = Table(seq1,seq2)
    aligner.needleman_Wunsch(sm,gap)
    print(aligner.recover_align())

    aligner.smith_Waterman(sm,gap)
    print(aligner.recover_align_local())

    print(needleman_Wunsch(seq1,seq2,sm.sm,gap))
    print(smith_Waterman(seq1,seq2,sm.sm,gap))
    pass